package com.baitiaojun.datasource.core.datasource;

import cn.hutool.core.map.MapUtil;
import com.alibaba.druid.pool.DruidDataSourceFactory;
import lombok.extern.slf4j.Slf4j;
import org.springframework.jdbc.datasource.lookup.AbstractRoutingDataSource;

import javax.sql.DataSource;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

@Slf4j
public class DynamicRoutingDataSource extends AbstractRoutingDataSource {

    /**
     * 用来判断ds是否存在
     */
    private static final Map<Object, Object> ds = new ConcurrentHashMap<>();

    /**
     * 每次打开回话连接, 都会去获取目标数据源名称, 默认获取master
     * @return
     */
    @Override
    protected Object determineCurrentLookupKey() {
        return DataSourceHolder.getDb();
    }

    /**
     * 首次启动设置默认数据源
     * @param targetDataSources
     */
    @Override
    public void setTargetDataSources(Map<Object, Object> targetDataSources) {
        super.setTargetDataSources(targetDataSources); //设置默认目标数据源
        super.afterPropertiesSet(); //初始化默认数据源
        ds.putAll(targetDataSources); //添加到数据源池
    }

    /**
     * 动态设置目标数据源
     * @param dbName
     * @param dataSource
     */
    public void setTargetDataSources(String dbName, DataSource dataSource) {
        super.setTargetDataSources(MapUtil.of(dbName, dataSource));
        super.afterPropertiesSet();
        ds.put(dbName, dataSource);
    }

    /**
     * 是否存在该数据源
     * @param dbName
     * @return
     */
    public boolean containDataSources(String dbName) {
        return ds.containsKey(dbName);
    }

    /**
     * 根据配置属性创建数据源
     * @param targetDataSources
     * @return
     */
    public DataSource createDataSource(Map<String, Object> targetDataSources) {
        DataSource dataSource;
        try {
            dataSource = DruidDataSourceFactory.createDataSource(targetDataSources);
        } catch (Exception e) {
            log.warn("无法获取指定数据源: {}", e.getMessage());
            throw new RuntimeException(e);
        }
        return dataSource;
    }
}
